Skip to content

WIP - removeBinary #414

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 9 commits into
base: feature/onnx-to-tosa
Choose a base branch
from
Open

WIP - removeBinary #414

wants to merge 9 commits into from

Conversation

@SushmitaThakallapalli1980
Copy link

@SushmitaThakallapalli1980 SushmitaThakallapalli1980 commented Aug 20, 2025 •
edited
Loading

REMOVE_BINARY summary:

This RemoveBinary class is a pattern-matching and rewrite utility that finds and removes binary ops (Add, Sub, Mul, Div) in quantized ONNX graphs when one input is effectively a constant. It folds the constant into the quantization parameters instead of keeping an explicit binary op in the graph.

This class rewrites patterns like:

 Dequant(x) ----\
                             Binary(Add/Sub/Mul/Div) → Quant → Dequant
Dequant(const)-/

into something like

  Dequant'(x, new_scale/zp) → Quant → Dequant
  
  Where the constant is folded into scale or zero-point. 
  
  Then, a check is performed to see if we can possibly remove the Quant → Dequant chain.

@tvivies-amd tvivies-amd requested a review from xiaohanAMD August 20, 2025 11:52
@tvivies-amd
Copy link

tvivies-amd commented Aug 20, 2025

I added @xiaohanAMD as reviewer too since he wrote the initial pattern and some improvements

tvivies-amd
tvivies-amd reviewed Aug 20, 2025
View reviewed changes
Copy link

@tvivies-amd tvivies-amd left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A couple of small comments, I like the usage of if constexpr and the state structure.

src/Dialect/ONNX/Transforms/DQBinaryQOpt.cpp Outdated
}
state.kValue = kValueOpt.value();

// Debug - To be removed
Copy link

@tvivies-amd tvivies-amd Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally you would not remove the debug ouptut that you used during the implementation, I would prefer if you could log them instead. @ehsan-toosi do you know if there is any logging capabilities in this repo ?

src/Dialect/ONNX/Transforms/DQBinaryQOpt.cpp
Comment on lines +225 to +226
if (failed(match_qdq(state, dqOp1, dqOp2)))
return failure();
Copy link

@tvivies-amd tvivies-amd Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: add bracket here too, in order to keep the coding style consistent

Copy link
Author

@SushmitaThakallapalli1980 SushmitaThakallapalli1980 Aug 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok

src/Dialect/ONNX/Transforms/DQBinaryQOpt.cpp Outdated
Comment on lines 228 to 231
// Debug - To be removed
llvm::outs() << "B. SUCCESS\n";
llvm::outs() << "kValue = " << state.kValue << "\n";
printOnnxNodeName(binaryOp, "[RemoveBinary] matched");
Copy link

@tvivies-amd tvivies-amd Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can put this in the match_qdq method and return directly the output of the call to match_qdq

src/Dialect/ONNX/Transforms/DQBinaryQOpt.cpp Outdated
Comment on lines 243 to 252
for (Value res : op->getResults()) {
for (Operation *user : res.getUsers()) {
if (auto q = dyn_cast<ONNXQuantizeLinearOp>(user)) {
quantOutputOp = q;
break;
}
}
if (quantOutputOp)
break;
}
Copy link

@tvivies-amd tvivies-amd Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the binary op is a qdq operation, I would expect that the output of the binOp only has one user that is the QuantizeLinear operation. This is what the python implementation is expecting and it is fine if we have the same check here

Copy link
Author

@SushmitaThakallapalli1980 SushmitaThakallapalli1980 Aug 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok.

src/Dialect/ONNX/Transforms/DQBinaryQOpt.cpp
.add<FoldBinaryThroughQDQ<ONNXDivOp>, FoldBinaryThroughQDQ<ONNXSubOp>,
FoldBinaryThroughQDQ<ONNXMulOp>, FoldBinaryThroughQDQ<ONNXAddOp>>(
&getContext());
if (failed(applyPatternsAndFoldGreedily(function, std::move(patterns))))
Copy link

@tvivies-amd tvivies-amd Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we need to use the greedy approach here, since we are not creating any new ops that would needs to be visited by the pass and same for the modified ops they do not need to be matched again. See documentation of the walker configuration: https://mlir.llvm.org/docs/PatternRewriter/#walk-pattern-rewrite-driver
What do you think ?

Copy link
Author

@SushmitaThakallapalli1980 SushmitaThakallapalli1980 Aug 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tried replacing applyPatternsAndFoldGreedily with walkAndApplyPatterns. But the code is crashing for some test cases. Not sure why. Therefore, restored greedy driver again.

@xiaohanAMD
Copy link

xiaohanAMD commented Aug 22, 2025 •
edited
Loading

this pattern is very complicated when there is fork in the match chain. Please add test case to fix this scenario, here Q1 has a fork, we expect not to fold into DQ1, but in Q2:

        const ---\
                 |
                 v
Q1 ---> DQ1 -> BinOp -> Q2 -> DQ2
    |
    \-> something else

@SushmitaThakallapalli1980
Copy link
Author

SushmitaThakallapalli1980 commented Aug 22, 2025

this pattern is very complicated when there is fork in the match chain. Please add test case to fix this scenario, here Q1 has a fork, we expect not to fold into DQ1, but in Q2:

        const ---\
                 |
                 v
Q1 ---> DQ1 -> BinOp -> Q2 -> DQ2
    |
    \-> something else

@xiaohanAMD
In the following https://jira.xilinx.com/browse/AIESW-8092, as per the initial scoping, folding is expected only on DQ1...Is this not the case?

@xiaohanAMD
Copy link

xiaohanAMD commented Aug 22, 2025

this pattern is very complicated when there is fork in the match chain. Please add test case to fix this scenario, here Q1 has a fork, we expect not to fold into DQ1, but in Q2:

        const ---\
                 |
                 v
Q1 ---> DQ1 -> BinOp -> Q2 -> DQ2
    |
    \-> something else

@xiaohanAMD In the following https://jira.xilinx.com/browse/AIESW-8092, as per the initial scoping, folding is expected only on DQ1...Is this not the case?

I didn't see that, sure, let's implement the basic in this PR. More complicate case we can do later.

@jorickert jorickert closed this Aug 25, 2025
@jorickert jorickert reopened this Aug 25, 2025
@SushmitaThakallapalli1980 SushmitaThakallapalli1980 marked this pull request as ready for review August 25, 2025 17:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tvivies-amd tvivies-amd tvivies-amd left review comments

@xiaohanAMD xiaohanAMD Awaiting requested review from xiaohanAMD

@jorickert jorickert Awaiting requested review from jorickert

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@SushmitaThakallapalli1980 @tvivies-amd @xiaohanAMD @jorickert